#include <rclcpp/rclcpp.hpp>
#include <opencv2/opencv.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <cv_bridge/cv_bridge.h>
#include <chrono>
extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
}

class RTSPStreamer : public rclcpp::Node
{
public:
    RTSPStreamer() : Node("rtsp_publisher")
    {
        // 订阅 ROS2 图像话题
        image_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/mb/video_stream_raw", 3, std::bind(&RTSPStreamer::image_callback, this, std::placeholders::_1));
        image_sub2_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/mb/video_stream", 3, std::bind(&RTSPStreamer::image_callback2, this, std::placeholders::_1));

        // RTSP 流 URL
        // 可以通过参数传递，也可以硬编码
        this->declare_parameter<std::string>("rtsp_url", "rtsp://127.0.0.1:8554/test2");
        this->get_parameter("rtsp_url", rtsp_url_);
        
        this->declare_parameter<int>("img_width", 640);
        this->get_parameter("img_width", img_width_);
        
        this->declare_parameter<int>("img_height", 480);
        this->get_parameter("img_height", img_height_);
        
        this->declare_parameter<int>("bit_rate", 3000000);
        this->get_parameter("bit_rate", bit_rate_);

        // 初始化 FFmpeg
        avformat_network_init();
        //initialize_ffmpeg(img_width_, img_height_);
        
        //previous_time = std::chrono::steady_clock::now();
    }

    ~RTSPStreamer()
    {
        av_write_trailer(fmt_ctx);
        avcodec_free_context(&codec_ctx);
        avformat_free_context(fmt_ctx);
        avformat_network_deinit();
    }

private:
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub_;
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr image_sub2_;

    std::chrono::steady_clock::time_point previous_time;

    std::string rtsp_url_;
    int img_width_ = 640;
    int img_height_ = 480;
    int bit_rate_ = 3000000;
    bool ffmpeg_Initialized = false;

    AVFormatContext *fmt_ctx = nullptr;
    AVCodecContext *codec_ctx = nullptr;
    AVStream *video_stream = nullptr;
    AVFrame *frame = nullptr;
    SwsContext *sws_ctx = nullptr;
    int frame_index = 0;
    
    void initialize_ffmpeg(int width, int height)
    {
        //const char *output_url = "rtsp://127.0.0.1:8554/test2";
        const char *output_url = rtsp_url_.c_str();

        avformat_alloc_output_context2(&fmt_ctx, nullptr, "rtsp", output_url);
        if (!fmt_ctx) {
            RCLCPP_ERROR(this->get_logger(), "无法分配输出上下文");
            return;
        }

        // 使用H.265编码器(HEVC)
        AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
        if (!codec) {
            RCLCPP_ERROR(this->get_logger(), "无法找到 H264 编码器");
            return;
        }

        codec_ctx = avcodec_alloc_context3(codec);
        codec_ctx->width = width;
        codec_ctx->height = height;
        codec_ctx->time_base = {1, 30};
        codec_ctx->framerate = {30, 1};
        codec_ctx->gop_size = 12;
        codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
        codec_ctx->bit_rate = bit_rate_;

        //codec_ctx->rc_buffer_size = 0;
        //codec_ctx->delay = 0;
        //codec_ctx->thread_count = 1;
        //codec_ctx->max_b_frames = 0;

#if 0
        // 强制包含 VPS/SPS/PPS
        codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
#else
        if (fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
            codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
#endif

        AVDictionary *opts = nullptr;

        //preset的参数调节编码速度和质量的平衡。
        //tune的参数值指定片子的类型，是和视觉优化的参数，
        //zerolatency: 零延迟，用在需要非常低的延迟的情况下，比如电视电话会议的编码
        //H.264
        if (codec_ctx->codec_id == AV_CODEC_ID_H264) {
            //av_dict_set(&opts, "preset", "slow", 0);
            //av_dict_set(&opts, "preset", "superfast", 0);
            av_dict_set(&opts, "preset", "veryfast", 0);
            av_dict_set(&opts, "tune", "zerolatency", 0);
            //av_dict_set(&opts, "profile", "main", 0);
        }

        //H.265
        if (codec_ctx->codec_id == AV_CODEC_ID_H265) {
            av_dict_set(&opts, "preset", "ultrafast", 0);
            av_dict_set(&opts, "tune", "zero-latency", 0);
            av_dict_set(&opts, "x265-params", "repeat-headers=1", 0);
        } 

        if (avcodec_open2(codec_ctx, codec, &opts) < 0) {
            RCLCPP_ERROR(this->get_logger(), "无法打开编码器");
            return;
        }

        video_stream = avformat_new_stream(fmt_ctx, codec);
        if (!video_stream) {
            RCLCPP_ERROR(this->get_logger(), "无法创建新的流");
            return;
        }

        video_stream->time_base = codec_ctx->time_base;

        avcodec_parameters_from_context(video_stream->codecpar, codec_ctx);

        if (!(fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
            if (avio_open(&fmt_ctx->pb, output_url, AVIO_FLAG_WRITE) < 0) {
                RCLCPP_ERROR(this->get_logger(), "无法打开输出 URL");
                return;
            }
        }

#if 1
        AVDictionary *opts2 = nullptr;
        av_dict_set(&opts2, "rtsp_streamer", "udp", 0);
        
        int ret = avformat_write_header(fmt_ctx, &opts2);
#else
        int ret = avformat_write_header(fmt_ctx, nullptr);
#endif
        if (ret < 0) {
            char err_buf[128];
            av_strerror(ret, err_buf, sizeof(err_buf));
            RCLCPP_ERROR(this->get_logger(), "无法写入流头部信息： %s", err_buf);
            return;
        }

        frame = av_frame_alloc();
        frame->format = AV_PIX_FMT_YUV420P;
        frame->width = codec_ctx->width;
        frame->height = codec_ctx->height;
        av_frame_get_buffer(frame, 0);

        sws_ctx = sws_getContext(
            codec_ctx->width, codec_ctx->height, AV_PIX_FMT_BGR24,
            codec_ctx->width, codec_ctx->height, AV_PIX_FMT_YUV420P,
            SWS_BILINEAR, nullptr, nullptr, nullptr);
    }

    void image_callback(const sensor_msgs::msg::Image::SharedPtr msg)
    {
        if (!ffmpeg_Initialized) {
            initialize_ffmpeg(msg->width, msg->height);
            ffmpeg_Initialized = true;
        }

        std::chrono::steady_clock::time_point now = std::chrono::steady_clock::now();
        if (std::chrono::duration_cast<std::chrono::milliseconds>(now - previous_time).count() < 3000) {
            return;
        }

        // 使用 cv_bridge 将 ROS 图像消息转换为 OpenCV 格式
        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            RCLCPP_ERROR(this->get_logger(), "cv_bridge 转换失败: %s", e.what());
            return;
        }

        if (cv_ptr->image.empty()) {
            return;
        }

        push_frame_to_rtsp(cv_ptr->image);
    }
    
    void image_callback2(const sensor_msgs::msg::Image::SharedPtr msg)
    {
        if (!ffmpeg_Initialized) {
            initialize_ffmpeg(msg->width, msg->height);
            ffmpeg_Initialized = true;
        }

        previous_time = std::chrono::steady_clock::now();

        // 使用 cv_bridge 将 ROS 图像消息转换为 OpenCV 格式
        cv_bridge::CvImagePtr cv_ptr;
        try {
            cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        } catch (cv_bridge::Exception& e) {
            RCLCPP_ERROR(this->get_logger(), "cv_bridge 转换失败: %s", e.what());
            return;
        }

        if (cv_ptr->image.empty()) {
            return;
        }

        push_frame_to_rtsp(cv_ptr->image);
    }
    
    void push_frame_to_rtsp(const cv::Mat& frame_cv)
    {
        cv::Mat frame_bgr = frame_cv.clone();
        if (frame_bgr.empty()) return;

        // 调整尺寸（可选）
        //cv::resize(frame_bgr, frame_bgr, cv::Size(img_width_, img_height_));

        const int in_linesize[1] = { static_cast<int>(frame_bgr.step[0]) };
        sws_scale(sws_ctx, &frame_bgr.data, in_linesize, 0, codec_ctx->height,
                  frame->data, frame->linesize);

        // 处理 pts 时间戳
        frame->pts = frame_index++;

        int ret = avcodec_send_frame(codec_ctx, frame);
        if (ret < 0) {
            RCLCPP_ERROR(this->get_logger(), "无法发送帧到编码器");
            return;
        }

        AVPacket pkt;
        av_init_packet(&pkt);
        pkt.data = nullptr;
        pkt.size = 0;

        ret = avcodec_receive_packet(codec_ctx, &pkt);
        if (ret == 0) {
            av_packet_rescale_ts(&pkt, codec_ctx->time_base, video_stream->time_base);
            
            //if (pkt.duration == 0) {
            //    pkt.duration = av_rescale_q(1, codec_ctx->time_base, video_stream->time_base);
            //}

            pkt.stream_index = video_stream->index;

            //std::cout << "PTS: " << pkt.pts << ", DTS: " << pkt.dts << ", Duration: " << pkt.duration << std:: endl;
            //std::cout << "encoder: " << codec_ctx->codec_id << ", AV_CODEC_ID_H264:" << AV_CODEC_ID_H264 << ", AV_CODEC_ID_HEVC:" << AV_CODEC_ID_HEVC << std:: endl;

            av_interleaved_write_frame(fmt_ctx, &pkt);
            av_packet_unref(&pkt);
        } else {
            RCLCPP_ERROR(this->get_logger(), "无法接收编码包");
        }
    }
};

int main(int argc, char *argv[])
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<RTSPStreamer>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}
